import math
import torch
from torch import nn
from d2l import torch as d2l


# @save
def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # X:3D张量，valid_lens:1D或2D张量
    if valid_lens is None:
        return nn.functional.softmax(X, dim=-1)
    else:
        shape = X.shape
        if valid_lens.dim() == 1:
            # 将 valid_lens 中的元素按顺序重复n次组成一个更大的1D张量
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            # 将 valid_lens 拉成一个1D的张量，保证元素数和X的元素数一样
            valid_lens = valid_lens.reshape(-1)
        # 最后一轴上被掩蔽的元素使用一个非常大的负值替换，从而其softmax输出为0
        X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens,
                              value=-1e6)
        return nn.functional.softmax(X.reshape(shape), dim=-1)


# 遮蔽注意力最大的问题是，在某些情况下，我们并不清楚哪些元素是需要参与到运算中来的。
# 以上方法专门用于处理 3维的特征信息
# 支持两种形式：
# 1. 1D张量，指的是在第一个维度上处理遮蔽信息。
# 2. 2D张量，指的是在前两个维度上处理遮蔽信息。
# 举例子：
# batch_size = 32 , seq_len = 100, num_hiddens = 10
# mask 的形状只支持 32 * 1 或者 32 * 100
# 如果是 batch_size * seq_len 的形式，则说明是将遮盖应用到每一个序列元素的特征上。
X = torch.randn([1, 5, 5])
X = torch.FloatTensor(
    [[[1., 1., 1., 1., 1.], # 只有第一个元素起作用
      [1., 1., 1., 1., 1.], # 只有第一个和第二个元素起作用
      [1., 1., 1., 1., 1.], # 以此类推
      [1., 1., 1., 1., 1.],
      [1., 1., 1., 1., 1.]]]
)
mask = torch.arange(
    1, 6, device=X.device).repeat(1, 1)
print(X)
attention_weights = masked_softmax(X, mask)
print(attention_weights)
x1 = torch.FloatTensor([[[1.0],
                         [1.0],
                         [1.0],
                         [1.0],
                         [1.0]]])
x1 = torch.randn([1, 5, 2])
print(x1)
print(torch.bmm(attention_weights, x1))
